package cn.doitedu;

import org.apache.hadoop.hive.ql.exec.UDAF;
import org.apache.hadoop.hive.ql.exec.UDAFEvaluator;
import org.roaringbitmap.longlong.Roaring64Bitmap;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.nio.ByteBuffer;

public class GuidAggBitMapUDAF extends UDAF {

    public static class GuidAggBitMapUDAFEvaluator implements UDAFEvaluator {

        Acc acc;

        public GuidAggBitMapUDAFEvaluator() {
            init();
        }

        // 初始化 中间聚合对象
        public void init() {

            // 构造一个空的bitmap
            Roaring64Bitmap bitmap = Roaring64Bitmap.bitmapOf();

            // 然后把这个bitmap序列化，赋值给 中间状态
            try {
                this.acc = new Acc(serBm(bitmap));

            } catch (IOException e) {
                throw new RuntimeException(e);
            }


        }

        // 对一条输入数据的聚合逻辑
        public boolean iterate(Long guid) {

            try {
                // 先反序列化“累加器bitmap"
                Roaring64Bitmap bitmap = deSerBm(this.acc.getBmBytes());
                bitmap.add(guid);

                // 把bitmap反序列化，更新到“累加器”
                this.acc.setBmBytes(serBm(bitmap));
                return true;

            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }


        // map端聚合完成后，输出结果
        public Acc terminatePartial() {
            return this.acc;
        }

        // reduce端聚合 map端输出的 结果
        public boolean merge(Acc otherAcc) {


            try {
                // 把otherAcc中的bm字节，反序列化成为bitmap对象
                Roaring64Bitmap otherBitmap = deSerBm(otherAcc.getBmBytes());

                // 把自己的acc中的bm字节，反序列化成为bitmap对象
                Roaring64Bitmap thisBitmap = deSerBm(this.acc.getBmBytes());


                // 然后，合并两个bitmap对象
                thisBitmap.or(otherBitmap);

                // 然后，将合并好的bitmap对象，序列化成字节，更新到“累加器”
                acc.setBmBytes(serBm(thisBitmap));


                return true;

            } catch (IOException e) {
                throw new RuntimeException(e);
            }

        }

        // reduce端，输出最终结果
        public byte[] terminate() {

            return acc.bmBytes;
        }
    }



    public static class Acc implements Serializable {

        private byte[] bmBytes;

        public Acc() {
        }

        public Acc(byte[] bmBytes) {
            this.bmBytes = bmBytes;
        }

        public byte[] getBmBytes() {
            return bmBytes;
        }

        public void setBmBytes(byte[] bmBytes) {
            this.bmBytes = bmBytes;
        }
    }


    public static byte[] serBm(Roaring64Bitmap bitmap) throws IOException {

        ByteArrayOutputStream bout = new ByteArrayOutputStream();
        DataOutputStream dout = new DataOutputStream(bout);
        bitmap.serialize(dout);

        return bout.toByteArray();


    }



    public static Roaring64Bitmap deSerBm(byte[] bmBytes) throws IOException {

        Roaring64Bitmap bitmap = Roaring64Bitmap.bitmapOf();
        bitmap.deserialize(ByteBuffer.wrap(bmBytes));

        return bitmap;
    }


}
